/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.model;

import java.util.HashMap;
import java.util.Map;

import net.javlov.Option;
import net.javlov.State;
import net.javlov.StateOptionPair;

@SuppressWarnings("rawtypes")
public class TransitionModel {

	/**
	 * The "regular" transition model. Maps as follows: (s,o) --> s' --> Transition
	 */
	protected Map<StateOptionPair, Map<State, Transition>> model;
	
	/**
	 * "Reverse" model, stores incoming transitions to states. Maps as
	 * <I>s'</I> --> Transition --> <I>s</I>. Note that we cannot use <I>s</I> as key here since there might be multiple incoming
	 * transitions from the same <I>s</I>.
	 * 
	 * Storing the reverse model is useful for e.g. prioritized sweeping.
	 */
	protected Map<State, Map<Transition, State>> revmodel;
	
	public TransitionModel() {
		model = new HashMap<StateOptionPair, Map<State, Transition>>();
		revmodel = new HashMap<State, Map<Transition, State>>();
	}
	
	public TransitionModel(int numStateOptions) {
		model = new HashMap<StateOptionPair, Map<State, Transition>>( (int) (numStateOptions / 0.75) );
		revmodel = new HashMap<State, Map<Transition, State>>( (int) (numStateOptions / 0.75) );
	}
	
	/**
	 * Sets the transition from (s,o) to sprime to probability p.
	 * @param s
	 * @param o
	 * @param sprime
	 * @param p
	 */
	public void add(State s, Option o, State sprime, double p) {
		Transition tr = getTransition(s, o, sprime);
		tr.setProb(p);
	}
	

	/**
	 * Increments the transition count from (s,o) to sprime by 1. Useful when learning a model.
	 * @param s
	 * @param o
	 * @param sprime
	 */
	public void count(State s, Option o, State sprime) {
		Transition tr = getTransition(s, o, sprime);
		tr.incrementCount();
	}
	
	public int getCount(State s, Option o, State sprime) {
		return getTransition(s, o, sprime).count;
	}
	
	/**
	 * 
	 * @param sprime
	 * @return a Map containing Transitions leading to sprime. The Value of the Map is the <I>source</I> state of each Transition.
	 */
	public Map<Transition, State> getIncomingTransitions(State sprime) {
		Map<Transition, State> revTransitions = revmodel.get(sprime);
		if ( revTransitions == null ) {
			revTransitions = new HashMap<Transition, State>();
			revmodel.put(sprime, revTransitions);
		}
		return revTransitions;
	}
	
	public double getProbability(State s, Option o, State sprime) {
		return getTransition(s, o, sprime).prob;
	}
	
	/**
	 * 
	 * @param s
	 * @param o
	 * @return a Map containing all transitions for this (s,o) pair. If no transitions are known, the map will be empty.
	 */
	public Map<State, Transition> getTransitions(State s, Option o) {
		StateOptionPair so = new StateOptionPair(s,o);
		Map<State, Transition> transitions = model.get(so);
		
		if ( transitions == null ) { //first time we see this pair
			transitions = new HashMap<State, Transition>();
			model.put(so, transitions);
		}
		return transitions;
	}
		
	
	/**
	 * Gets the transition (s,o,s'). Creates a new one if it does not exist; in this case, also adds an entry to
	 * the reverse transition model.
	 * 
	 * @param s
	 * @param o
	 * @param sprime
	 * @return
	 */
	public Transition getTransition(State s, Option o, State sprime) {	
		//get all transitions for the state-option pair (s,o)
		Map<State, Transition> transitions = getTransitions(s, o);
		
		//get specific transition to this sprime
		Transition tr = transitions.get(sprime);
		
		//if the transition is unknown, create a new one and add an entry to the reverse model
		if ( tr == null ) {
			tr = new Transition(s,o,sprime);
			transitions.put(s, tr);
			
			Map<Transition, State> revTransitions = getIncomingTransitions(sprime);
			revTransitions.put(tr, s);
		}
		return tr;
	}

	/**
	 * Updates probabilities from counts (note that this will overwrite any previous probabilities, also
	 * ones set using {@link #add(State, Option, State, double)}).
	 */
	public void update() {
		double sum;
		for ( Map<State, Transition> transitions : model.values() ) {
			sum = 0;
			for ( Transition tr : transitions.values() )
				sum += tr.count;
			for ( Transition tr : transitions.values() )
				tr.setProb(tr.count / sum);		
		}
	}
	
	
	public class Transition {
		State s, sprime;
		Option o;
		double prob;
		int count;
		
		Transition(State s, Option o, State sprime, double p) {
			this.s = s;
			this.o = o;
			this.sprime = sprime;
			setProb(p);
			setCount(0);
		}
		
		Transition(State s, Option o, State sprime) {
			this.s = s;
			this.o = o;
			this.sprime = sprime;
			setProb(0);
			setCount(0);
		}
		
		public int getCount() {
			return count;
		}
		
		public double getProbability() {
			return prob;
		}
		
		public Option getOption() {
			return o;
		}
		
		public State getSourceState() {
			return s;
		}
		
		public State getTargetState() {
			return sprime;
		}
		
		void incrementCount() {
			count++;
		}
		
		void setCount(int c) {
			count = c;
		}
		
		void setProb(double p) {
			prob = p;
		}
		
		@Override
		public boolean equals(Object obj) {
			if ( obj == null || !(obj instanceof Transition) )
				return false;
			Transition tr = (Transition)obj;
			return tr.s.equals(s) && tr.o.equals(o) && tr.sprime.equals(sprime);
		}
		
		@Override
		public int hashCode() {
			return 31*s.hashCode() + 31*o.hashCode() + sprime.hashCode();
		}
	}
	
}
